package com.wenyl.tree;

import java.util.Stack;

import lombok.Data;

/**
 * AVL树的实现类,注意高度是指节点到其最深的叶子节点的最长路径的长度
 */
@Data
public class AVLTree {
    /**
     * 二叉搜索树的根节点
     */
    private AVLTNode root;

    /**
     * 向二叉搜索树中插入一个新的节点
     * @param data 要插入的数据
     */
    public void insert(int data){
        if (root == null) {
            root = new AVLTNode(data);
            return;
        }
        AVLTNode currentNode = root;
        Stack<AVLTNode> stack = new Stack<>();
        // 遍历树，找到插入位置
        while (currentNode != null) {
            stack.push(currentNode);
            if (data < currentNode.getData()) {
                if (currentNode.getLeftChild() == null) {
                    currentNode.setLeftChild(new AVLTNode(data));
                    break;
                }
                currentNode = currentNode.getLeftChild();
            } else if (data > currentNode.getData()) {
                if (currentNode.getRightChild() == null) {
                    currentNode.setRightChild(new AVLTNode(data));
                    break;
                }
                currentNode = currentNode.getRightChild();
            } else {
                System.out.println("数据重复，插入失败");
                return;
            }
        }

        // 更新插入时走过的节点高度，这里是从插入位置到根节点的路径上的所有节点，反向遍历
        while (!stack.isEmpty()) {
            currentNode = stack.pop();
            // 更新当前节点的高度
            updateHeight(currentNode);
            // 检查是否平衡,之类平衡的时当前节点，更新当前节点后，需要更新父节点的子节点
            currentNode = rebalance(currentNode);

            // 如果栈不为空，则更新父节点的子节点
            if(!stack.isEmpty()){
                AVLTNode parentNode = stack.peek();
                if(currentNode.getData() < parentNode.getData()){
                    parentNode.setLeftChild(currentNode);
                }else{
                    parentNode.setRightChild(currentNode);
                }
            }else{
                // 如果栈为空，则当前节点为根节点
                root = currentNode;
            }
        }
    }

    /**
     * 检查节点是否平衡
     * @param node 当前节点
     */
    private AVLTNode rebalance(AVLTNode node) {
        int balanceFactor = getBalanceFactor(node);
        
        // 左子树高度大于右子树高度，并且左子树的左子树高度大于等于左子树的右子树高度，则进行右旋
        if (balanceFactor > 1 && getBalanceFactor(node.getLeftChild()) >= 0) {
            return rotateRight(node);
        }
        // 左子树高度大于右子树高度，并且左子树的右子树高度大于等于左子树的左子树高度，则进行左右旋
        if(balanceFactor>1 && getBalanceFactor(node.getLeftChild())<0){
            return rotateLeftRight(node);
        }
        // 右子树高度大于左子树高度，并且右子树的左子树高度大于等于右子树的右子树高度，则进行左旋
        if(balanceFactor<-1 && getBalanceFactor(node.getRightChild())<=0){
            return rotateLeft(node);
        }
        // 右子树高度大于左子树高度，并且右子树的右子树高度大于等于右子树的左子树高度，则进行右左旋
        if(balanceFactor<-1 && getBalanceFactor(node.getRightChild())>0){
            return rotateRightLeft(node);
        }
        // 平衡，返回当前节点
        return node;
    }

    /**
     * 获取节点的平衡因子，左子树的高度减去右子树的高度
     * @param node 当前节点
     * @return 平衡因子
     */
    private int getBalanceFactor(AVLTNode node){
        return getHeight(node.getLeftChild()) - getHeight(node.getRightChild());
    }

    /**
     * 获取节点的高度,若节点不存在返回-1
     * @param node 当前节点
     * @return 节点的高度
     */
    private int getHeight(AVLTNode node){
        return node == null ? -1 : node.getNodeHeight();
    }

    /**
     * 更新节点的高度
     * @param node 当前节点
     */
    private void updateHeight(AVLTNode node) {
        // 如果当前节点为空，则返回
        if (node == null) {
            return;
        }
        // 更新当前节点的高度，从该节点到其最深的叶子节点的最长路径的长度
        node.setNodeHeight(Math.max(getHeight(node.getLeftChild()), getHeight(node.getRightChild())) + 1);
    }



    /**
     * 从二叉搜索树中删除指定数据的节点
     * @param data 要删除的数据
     */
    public void delete(int data){
        if (root == null) {
            System.out.println("树为空，删除失败");
            return;
        }
        AVLTNode currentNode = root;
        AVLTNode parentNode = null;
        Stack<AVLTNode> stack = new Stack<>();
        // 遍历树，找到要删除的节点,栈中不包含要删除的节点
        while (currentNode != null && currentNode.getData() != data) {
            stack.push(currentNode);
            parentNode = currentNode;
            if (data < currentNode.getData()) {
                currentNode = currentNode.getLeftChild();
            } else {
                currentNode = currentNode.getRightChild();
            } 
        }    
        if (currentNode == null) {
            System.out.println("数据不存在，删除失败");
            return;
        }
        if(currentNode.getLeftChild() == null && currentNode.getRightChild() == null){
            deleteNodeWithOneOrZeroChild(currentNode, parentNode);
        }else if(currentNode.getLeftChild() == null && currentNode.getRightChild() != null){
            deleteNodeWithOneOrZeroChild(currentNode, parentNode);
        }else if(currentNode.getLeftChild() != null && currentNode.getRightChild() == null){
            deleteNodeWithOneOrZeroChild(currentNode, parentNode);
        }else{
            deleteNodeWithTwoNodes(currentNode, parentNode);
        }
        
        while(!stack.isEmpty()){
            currentNode = stack.pop();
            updateHeight(currentNode);
            currentNode = rebalance(currentNode);
            if (!stack.isEmpty()) {
                parentNode = stack.peek();
                if (currentNode.getData() < parentNode.getData()) {
                    parentNode.setLeftChild(currentNode);
                } else {
                    parentNode.setRightChild(currentNode);
                }
            } else {
                root = currentNode;
            }
        }
    }

    /**
     * 删除只有一个或没有子节点的节点
     * @param currentNode 当前节点
     * @param parentNode 父节点
     * @return 删除的节点
     */
    private AVLTNode deleteNodeWithOneOrZeroChild(AVLTNode currentNode, AVLTNode parentNode){
        AVLTNode child = currentNode.getLeftChild() == null ? currentNode.getRightChild() : currentNode.getLeftChild();
        if(parentNode == null){
            root = child;
        }else if(parentNode.getLeftChild() == currentNode){
            parentNode.setLeftChild(child);
        }else{
            parentNode.setRightChild(child);
        }
        currentNode.setLeftChild(null);
        currentNode.setRightChild(null);
        updateHeight(parentNode);
        return child;
    }
    /**
     * 删除左右子树都有的节点
     * @param currentNode 当前节点
     * @param parentNode 父节点
     * @return 删除的节点
     */
    private AVLTNode deleteNodeWithTwoNodes(AVLTNode currentNode, AVLTNode parentNode){
        AVLTNode rightMinNode = findRightMinNode(currentNode);
        rightMinNode.setLeftChild(currentNode.getLeftChild());
        // 如果右子树最小节点不是当前节点的右子树，则将右子树最小节点的右子树设置为当前节点的右子树
        if(rightMinNode != currentNode.getRightChild()){
            rightMinNode.setRightChild(currentNode.getRightChild());
        }
        // 如果父节点为空，说明要删除的节点是根节点
        if(parentNode == null){
            root = rightMinNode;
        }else if(parentNode.getLeftChild() == currentNode){
            parentNode.setLeftChild(rightMinNode);
        }else{
            parentNode.setRightChild(rightMinNode);
        }
        // 将当前节点从树中移除,等待垃圾回收
        currentNode.setLeftChild(null);
        currentNode.setRightChild(null);
        updateHeight(rightMinNode);
        return rightMinNode;
    }

    // todo 处理右子树最小节点的右子树的情况
    /**
     * 查找右子树中的最小节点
     * @param currentNode 当前节点
     * @return 右子树中的最小节点
     */
    private AVLTNode findRightMinNode(AVLTNode currentNode) {
        AVLTNode minNode = currentNode.getRightChild();
        AVLTNode parentNode = currentNode;
        
        while (minNode.getLeftChild() != null) {
            parentNode = minNode;
            minNode = minNode.getLeftChild();
        }
        // 将最小节点从右子树中移除,用来替换被删除的节点
        // 若最小节点就是删除节点的右子树，则不做处理
        if(parentNode != currentNode){
            parentNode.setLeftChild(minNode.getRightChild());
        }

        // 更新最小节点的高度
        updateHeight(parentNode);
        return minNode;
    }

    /**
     * 前序遍历二叉树
     * 遍历顺序：根节点 -> 左子树 -> 右子树
     * @param currentNode 当前节点
     */
    public void preOrderTraversal(AVLTNode currentNode){
        if(currentNode == null){
            return;
        }
        System.out.print(currentNode.getData()+" ");
        preOrderTraversal(currentNode.getLeftChild());
        preOrderTraversal(currentNode.getRightChild());
    }

    /**
     * 中序遍历二叉树
     * 遍历顺序：左子树 -> 根节点 -> 右子树
     * @param currentNode 当前节点
     */
    public void inOrderTraversal(AVLTNode currentNode){
        if(currentNode == null){
            return;
        }
        inOrderTraversal(currentNode.getLeftChild());
        System.out.print(currentNode.getData()+" ");
        inOrderTraversal(currentNode.getRightChild());
    }

    /**
     * 后序遍历二叉树
     * 遍历顺序：左子树 -> 右子树 -> 根节点
     * @param currentNode 当前节点
     */
    public void postOrderTraversal(AVLTNode currentNode){
        if(currentNode == null){
            return;
        }
        postOrderTraversal(currentNode.getLeftChild());
        postOrderTraversal(currentNode.getRightChild());
        System.out.print(currentNode.getData()+" ");
    }

    /**
     * LL形，做右旋，将当前节点设置为左孩子的右节点，如果左孩子右节点不为空，则将左孩子的右节点设置为当前节点的左节点
     * @param node 当前节点
     * @return 旋转后的新节点
     */
    public AVLTNode rotateRight(AVLTNode node){
        AVLTNode leftChild = node.getLeftChild();
        // 将左子树的右子树设置为当前节点的左子树
        node.setLeftChild(leftChild.getRightChild());
        // 将当前节点设置为左子树的右子树
        leftChild.setRightChild(node);

        updateHeight(node);
        updateHeight(leftChild);

        return leftChild;
    }

    /**
     * RR形，做左旋，将当前节点设置为右孩子的左节点，如果右孩子左节点不为空，则将右孩子的左节点设置为当前节点的右节点
     * @param node 当前节点
     * @return 旋转后的新节点
     */
    public AVLTNode rotateLeft(AVLTNode node){
        AVLTNode rightChild = node.getRightChild();
        // 将右子树的左子树设置为当前节点的右子树
        node.setRightChild(rightChild.getLeftChild());
        // 将当前节点设置为右子树的左子树
        rightChild.setLeftChild(node);

        updateHeight(node);
        updateHeight(rightChild);

        return rightChild;
    }

    /**
     * LR形，先对左孩子做左旋，再对当前节点做右旋
     * @param node 当前节点
     * @return 旋转后的新节点
     */
    public AVLTNode rotateLeftRight(AVLTNode node){
        node.setLeftChild(rotateLeft(node.getLeftChild()));
        return rotateRight(node);
    }

    /**
     * RL形，先对右孩子做右旋，再对当前节点做左旋
     * @param node 当前节点
     * @return 旋转后的新节点
     */
    public AVLTNode rotateRightLeft(AVLTNode node){
        node.setRightChild(rotateRight(node.getRightChild()));
        return rotateLeft(node);
    }   
}

/**
 * AVL树的节点类
 */
@Data
class AVLTNode{
    /**
     * 节点存储的整数数据
     */
    private int data;
    
    /**
     * 左子节点
     */
    private AVLTNode leftChild;
    
    /**
     * 右子节点
     */
    private AVLTNode rightChild;
    
    /**
     * 节点的高度，从该节点到其最深的叶子节点的最长路径的长度
     */
    private int nodeHeight;
    
    /**
     * 构造一个新的AVLTNode
     * @param data 节点存储的整数数据
     */
    public AVLTNode(int data){
        this.data = data;
        this.nodeHeight = 0;
    }
}
